# -*- coding:utf-8 -*-

import torch
import torch.nn as nn
# 进行句子的截断补齐（规范长度）
from keras.preprocessing import sequence
import pandas as pd
from collections import Counter
from functools import reduce
from sklearn.utils import shuffle
import torch.optim as optim
from net import Net

# 预训练模型来源
source = 'huggingface/pytorch-transformers'

# 加载的预训练模型的名字
model_name = 'bert-base-chinese'

# 选定加载模型的哪⼀部分,这⾥是模型的映射器
part = 'tokenizer'
tokenizer = torch.hub.load(source, part, model_name)

# # 加载不带头的预训练模型
part = 'model'
model = torch.hub.load(source, part, model_name)

net = Net()

# 定义交叉熵损失函数
criterion = nn.CrossEntropyLoss()

# 定义SGD优化方法
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)


def get_bert_encode(text, cutlen=32):
    # 使用bert-chinese编码中文文本
    # text 要进行编码的文本
    # return 使用bert编码后的文本张量表示
    # 首先使用字符映射器对每个汉字进行映射
    # 这里需要注意，bert的tokenizer映射后会为结果前后添加开始和结束标记即101和102
    # 这对于多段文本的编码是有意义的，但在我们这里没意义，因此使用[1:-1]对头和尾进行切片
    indexed_tokens = tokenizer.encode(text[:cutlen])[1:-1]
    # 对映射后的句子进行截断补齐
    indexed_tokens = sequence.pad_sequences([indexed_tokens], cutlen)
    # 之后将列表结构转化为tensor
    tokens_tensor = torch.LongTensor(indexed_tokens)

    # 使用模型不自动计算梯度
    with torch.no_grad():
        # 调用模型获得隐层输出
        model_result = model(tokens_tensor)
        encoded_layers = model_result.last_hidden_state
    # 输出的隐层是一个三维张量，最外层一维是1，我们使用【0】降去它
    encoded_layers = encoded_layers[0]
    return encoded_layers


def data_loader(train_data_path, valid_data_path, batch_size):
    # 从持久化文件中加载数据
    # train_data_path:训练数据路径
    # valid_data_path:验证数据路径
    # batch_size：训练和验证数据集的批次大小
    # return 训练数据生成器，验证数据生成器，训练数据数量，验证数据数量
    # 使用pd进行csv数据的读取，并去除第一行的列名
    train_data = pd.read_csv(train_data_path, header=None, sep="\t").drop([0])
    valid_data = pd.read_csv(valid_data_path, header=None, sep="\t").drop([0])

    # # 打印训练集和验证集上的正负样本数量
    # print("训练数据集的正负样本数量：")
    # print(dict(Counter(train_data[0].values)))
    # print("验证证数据集的正负样本数量：")
    # print(dict(Counter(valid_data[0].values)))

    # 验证数据集中的数据总数至少能够满足一个批次
    if len(valid_data) < batch_size:
        raise ("Batch size or split not match!")

    def _loader_generator(data):
        # 获取训练集/验证集的每个批次数据的生成器
        # data:训练数据或验证数据
        # return 一个批次的训练数据或验证数据的生成器

        # 以每个批次的间隔遍历数据集
        for batch in range(0, len(data), batch_size):
            # 定义batch数据的张量列表
            batch_encoded = []
            batch_labels = []
            # 将一个bitch_size大小的数据转换成列表形式，并进行逐条遍历
            for item in shuffle(data.values.tolist())[batch: batch + batch_size]:
                # 使用bert中文模型进行编码
                encoded = get_bert_encode(item[0])
                # 将编码后的每条数据装进预先定义好的列表中
                batch_encoded.append(encoded)
                # 同样将对应的该batch的标签装进labels列表中
                batch_labels.append([int(item[1])])
            # 使用reduce高阶函数将列表中的数据转换成模型需要的张量形式
            # encoded的形状是(batch_size, 2*max_len, embedding_size)
            encoded = reduce(lambda x, y: torch.cat((x, y), dim=0), batch_encoded)
            labels = torch.tensor(reduce(lambda x, y: x + y, batch_labels))
            # 以生成器的方式返回数据和标签
            yield (encoded, labels)

    # 对训练集和验证集分别使用_loader_generator函数，返回对应的生成器
    # 最后还要返回训练集和验证集的样本数量
    return _loader_generator(train_data), _loader_generator(valid_data), len(train_data), len(valid_data)


def train(train_data_labels):
    # 训练函数，在这个过程中将要更新模型参数，并收集准确率和损失
    # train_data_labels：训练数据和标签的生成器对象
    # return：整个训练过程的平均损失之和以及正确标签的累加数

    # 定义训练过程的初始化损失和准确率累加数
    train_running_loss = 0.0
    train_running_acc = 0.0
    # 循环遍历训练数据和标签生成器，每个批次更新一次模型参数
    for train_tensor, train_labels in train_data_labels:
        # 初始化该批次的优化器
        optimizer.zero_grad()
        # 使用微调网络获取输出
        train_outputs = net(train_tensor)
        # 得到该批次下的平均损失
        train_loss = criterion(train_outputs, train_labels)
        # 将该批次的平均损失加到train_running_loss中
        train_running_loss += train_loss.item()
        # 损失反向传播
        train_loss.backward()
        # 优化器更新模型参数
        optimizer.step()
        # 将该批次中正确的标签数量进行累加，以便之后计算准确率
        train_running_acc += (train_outputs.argmax(1) == train_labels).sum().item()

    return train_running_loss, train_running_acc


def valid(valid_data_labels):
    # 训练函数，在这个过程中将要更新模型参数，并收集准确率和损失
    # valid_data_lables：验证数据和标签的生成器对象
    # return：整个训练过程的平均损失之和以及正确标签的累加数

    # 定义训练过程的初始化损失和准确率累加数
    valid_running_loss = 0.0
    valid_running_acc = 0.0
    # 循环遍历训练数据和标签生成器，每个批次更新一次模型参数
    for valid_tensor, valid_labels in valid_data_labels:
        # 不自动更新梯度
        with torch.no_grad():
            # 使用微调网络获取输出
            valid_outputs = net(valid_tensor)
            # 得到该批次下的平均损失
            valid_loss = criterion(valid_outputs, valid_labels)
            # 将该批次的平均损失加到valid_running_loss中
            valid_running_loss += valid_loss.item()
            # 将该批次中正确的标签数量进行累加，以便之后计算准确率
            valid_running_acc += (valid_outputs.argmax(1) == valid_labels).sum().item()

    return valid_running_loss, valid_running_acc

def label2msg(label):
    if label == 0:
        return "差评"
    else:
        return "好评"

if __name__ == "__main__":
    # text = "早餐不好，服务不到位，晚餐无西餐，早餐晚餐相同，房间条件不好"
    # encoded_layers = get_bert_encode(text)
    # print(encoded_layers)
    # print(encoded_layers.shape)

    train_data_path = "./cn_data/SST-2/train.tsv"
    valid_data_path = "./cn_data/SST-2/dev.tsv"

    # 定义训练轮数
    epochs = 4

    # 定义批次样本数量
    batch_size = 32

    # 保存路径
    MODEL_PATH = './BERT_net_'+str(epochs)+'_'+str(batch_size)+'.pth'

    # batch_size = 16
    # train_data_labels, valid_data_labels, \
    # train_data_len, valid_data_len = data_loader(train_data_path, valid_data_path, batch_size)
    #
    # print(next(train_data_labels))
    # print(next(valid_data_labels))
    # print("train_data_len:", train_data_len)
    # print("valid_data_len:", valid_data_len)

    #######################################训练###########################
    # # 进行指定轮次的训练
    # for epoch in range(epochs):
    #     # 打印轮次
    #     print("Epoch:", epoch + 1)
    #
    #     # 通过数据加载器获得训练数据和验证数据生成器，以及对应的样本数量
    #     train_data_labels, valid_data_labels, \
    #         train_data_len, valid_data_len = data_loader(train_data_path, valid_data_path, batch_size)
    #     # 调用训练函数进行训练
    #     train_running_loss, train_running_acc = train(train_data_labels)
    #     # 调用验证函数进行验证
    #     valid_running_loss, valid_running_acc = valid(valid_data_labels)
    #     # 计算每一轮的平均损失，train_running_loss和valid_running_loss是每个批次的平均损失之和
    #     # 因此将它们乘以batch_size就得到了该轮的总损失，除以样本数即该轮次的平均损失
    #     train_average_loss = train_running_loss * batch_size / train_data_len
    #     valid_average_loss = valid_running_loss * batch_size / valid_data_len
    #
    #     # train_running_acc 和 valid_running_acc是每个批次的正确标签累加和
    #     # 因此只需除以对应样本总数即是该轮次的准确率
    #     train_average_acc = train_running_acc / train_data_len
    #     valid_average_acc = valid_running_acc / valid_data_len
    #
    #     # 打印训练损失和准确率以及验证损失和准确率
    #     print("Train Loss:", train_average_loss, "|", "Train Acc:", train_average_acc)
    #     print("Valid Loss:", valid_average_loss, "|", "Valid Acc:", valid_average_acc)
    #
    # print("Finished Training")
    #
    # # 保存模型参数
    # torch.save(net.state_dict(), MODEL_PATH)
    # print("Finished Saving")

    #######################################训练###########################

    #######################################预测###########################
    # 加载模型参数
    net.load_state_dict(torch.load(MODEL_PATH))

    texts = [
        "酒店环境很好 交通很方便 附近有比较多吃东西的地方 我是晚上到的 12点多在附近还是可以很方便吃到东西",
        "酒店环境很好 12点多在附近还是可以很方便吃到东西",
        "外观旧,停车场小,乱,还另外收钱,不过酒店还干净.",
        "乱,还另外收钱,不过酒店还干净.",
        "房间有点旧",
        "服务质量都只能用有差来形容",
        "设施稍微有点旧但是可以接受，但是606的价格还不含早餐有点高了",
        "酒店可真是太差了",
        "房价贵了点",
        "太吵,洗手间没有热水",
        "房间太旧",
        "位置比较偏打不车不方便"
    ]

    for text in texts:
        print("输入文本为：", text)
        with torch.no_grad():
            output = net(get_bert_encode(text))
            # 从output中取出最大值对应的索引
            print("预测标签为：", label2msg(torch.argmax(output).item()))

    #######################################预测###########################
